#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
高斯数据库表字段批量添加脚本（优化版）
功能：读取columns.txt文件，生成批量ALTER TABLE语句
特点：
1. 支持批量执行模式
2. 添加错误处理
3. 支持字段长度自定义
4. 生成执行统计信息
作者：AI Assistant
"""

import os
import re
from collections import defaultdict

def map_data_type_advanced(original_type):
    """
    高级数据类型映射，支持更多类型和长度推断
    
    Args:
        original_type (str): 原始数据类型
        
    Returns:
        str: 映射后的数据类型
    """
    type_mapping = {
        'bigint': 'BIGINT',
        'string': 'VARCHAR(500)',  # 字符串类型默认长度
        'int4': 'INTEGER',
        'int': 'INTEGER',
        'integer': 'INTEGER',
        'varchar': 'VARCHAR(500)',
        'text': 'TEXT',
        'decimal': 'DECIMAL(18,2)',
        'numeric': 'NUMERIC(18,2)',
        'float': 'FLOAT',
        'double': 'DOUBLE PRECISION',
        'date': 'DATE',
        'datetime': 'TIMESTAMP',
        'timestamp': 'TIMESTAMP',
        'boolean': 'BOOLEAN',
        'bool': 'BOOLEAN'
    }
    
    # 处理带长度的varchar类型
    if original_type.lower().startswith('varchar('):
        return original_type.upper()
    
    # 处理decimal/numeric类型
    if original_type.lower().startswith(('decimal(', 'numeric(')):
        return original_type.upper()
    
    return type_mapping.get(original_type.lower(), 'VARCHAR(500)')

def analyze_column_types(columns):
    """
    分析字段类型分布
    
    Args:
        columns (list): 字段信息列表
        
    Returns:
        dict: 类型统计信息
    """
    type_stats = defaultdict(int)
    
    for _, _, data_type, _ in columns:
        mapped_type = map_data_type_advanced(data_type)
        type_stats[mapped_type] += 1
    
    return dict(type_stats)

def read_columns_file_enhanced(file_path):
    """
    增强版字段文件读取，支持更多格式和错误处理
    
    Args:
        file_path (str): columns.txt文件路径
        
    Returns:
        tuple: (字段列表, 错误行列表)
    """
    columns = []
    error_lines = []
    
    try:
        with open(file_path, 'r', encoding='utf-8') as f:
            for line_num, line in enumerate(f, 1):
                line = line.strip()
                if not line or line.startswith('#'):  # 跳过空行和注释行
                    continue
                    
                parts = line.split('\t')
                if len(parts) >= 4:
                    table_name = parts[0].strip()
                    column_name = parts[1].strip()
                    data_type = parts[2].strip()
                    comment = parts[3].strip()
                    
                    # 验证字段名（避免SQL注入和语法错误）
                    if re.match(r'^[A-Za-z_][A-Za-z0-9_]*$', column_name):
                        columns.append((table_name, column_name, data_type, comment))
                    else:
                        error_lines.append((line_num, f"无效字段名: {column_name}"))
                else:
                    error_lines.append((line_num, f"格式不正确: {line}"))
                    
    except FileNotFoundError:
        print(f"错误：找不到文件 {file_path}")
        return [], []
    except Exception as e:
        print(f"错误：读取文件时发生异常 {e}")
        return [], []
        
    return columns, error_lines

def generate_batch_alter_statements(columns, target_table_name, full_table_name, batch_size=50):
    """
    生成批量ALTER TABLE语句
    
    Args:
        columns (list): 字段信息列表
        target_table_name (str): 目标表名
        full_table_name (str): 完整表名
        batch_size (int): 每批处理的字段数量
        
    Returns:
        list: SQL语句列表
    """
    statements = []
    
    # 过滤出属于目标表的字段
    target_columns = [col for col in columns if col[0] == target_table_name]
    
    if not target_columns:
        print(f"警告：没有找到表 {target_table_name} 的字段定义")
        return statements
    
    print(f"找到 {len(target_columns)} 个字段需要添加到表 {target_table_name}")
    
    # 添加文件头注释
    statements.extend([
        "-- ======================================================",
        "-- 高斯数据库表字段批量添加脚本",
        "-- 自动生成的ALTER TABLE语句",
        f"-- 目标表: {full_table_name}",
        f"-- 字段总数: {len(target_columns)}",
        f"-- 生成时间: {__import__('datetime').datetime.now().strftime('%Y-%m-%d %H:%M:%S')}",
        "-- 执行前请仔细检查！",
        "-- ======================================================",
        "",
        "-- 开始事务",
        "BEGIN;",
        ""
    ])
    
    # 分批处理字段
    for i in range(0, len(target_columns), batch_size):
        batch = target_columns[i:i+batch_size]
        batch_num = i // batch_size + 1
        
        statements.append(f"-- 批次 {batch_num}: 字段 {i+1} 到 {min(i+batch_size, len(target_columns))}")
        statements.append("")
        
        for table_name, column_name, data_type, comment in batch:
            # 映射数据类型
            mapped_type = map_data_type_advanced(data_type)
            
            # 生成ALTER TABLE语句
            alter_statement = f"ALTER TABLE {full_table_name} ADD COLUMN {column_name} {mapped_type};"
            statements.append(alter_statement)
        
        statements.append("")
        
        # 添加注释语句
        statements.append(f"-- 批次 {batch_num} 字段注释")
        for table_name, column_name, data_type, comment in batch:
            # 转义单引号
            escaped_comment = comment.replace("'", "''")
            comment_statement = f"COMMENT ON COLUMN {full_table_name}.{column_name} IS '{escaped_comment}';"
            statements.append(comment_statement)
        
        statements.extend(["", f"-- 批次 {batch_num} 完成", ""])
    
    # 添加文件尾
    statements.extend([
        "-- 提交事务",
        "COMMIT;",
        "",
        "-- 脚本执行完成",
        f"-- 总共添加了 {len(target_columns)} 个字段"
    ])
    
    return statements

def save_statements_with_summary(statements, output_file, type_stats, error_lines):
    """
    保存SQL语句并生成执行摘要
    
    Args:
        statements (list): SQL语句列表
        output_file (str): 输出文件路径
        type_stats (dict): 类型统计信息
        error_lines (list): 错误行信息
    """
    try:
        with open(output_file, 'w', encoding='utf-8') as f:
            for statement in statements:
                f.write(statement + "\n")
                
        print(f"SQL语句已保存到文件: {output_file}")
        
        # 生成摘要文件
        summary_file = output_file.replace('.sql', '_summary.txt')
        with open(summary_file, 'w', encoding='utf-8') as f:
            f.write("字段添加脚本执行摘要\n")
            f.write("=" * 50 + "\n\n")
            
            f.write("数据类型分布:\n")
            for data_type, count in sorted(type_stats.items()):
                f.write(f"  {data_type}: {count} 个字段\n")
            
            if error_lines:
                f.write("\n处理错误:\n")
                for line_num, error in error_lines:
                    f.write(f"  第{line_num}行: {error}\n")
            else:
                f.write("\n✓ 所有字段处理成功，无错误\n")
            
            f.write(f"\n总计: {sum(type_stats.values())} 个字段\n")
            f.write(f"SQL文件: {os.path.basename(output_file)}\n")
        
        print(f"执行摘要已保存到: {summary_file}")
        
    except Exception as e:
        print(f"错误：保存文件时发生异常 {e}")

def main():
    """
    主函数
    """
    # 文件路径配置
    current_dir = os.path.dirname(os.path.abspath(__file__))
    columns_file = os.path.join(current_dir, 'columns.txt')
    output_file = os.path.join(current_dir, 'batch_alter_table_statements.sql')
    
    # 目标表名配置
    target_table_name = 'exec_imp_tag_bscom_bkupb_intg_upf_basic_info_np'
    full_table_name = 'market.exec_imp_tag_bscom_bkupb_intg_upf_basic_info_np'
    
    print("高斯数据库表字段批量添加脚本（优化版）")
    print("=" * 50)
    print(f"输入文件: {columns_file}")
    print(f"输出文件: {output_file}")
    print(f"目标表名: {target_table_name}")
    print(f"完整表名: {full_table_name}")
    print("-" * 50)
    
    # 读取字段定义
    columns, error_lines = read_columns_file_enhanced(columns_file)
    
    if not columns:
        print("没有读取到任何有效字段定义，程序退出")
        return
    
    print(f"总共读取到 {len(columns)} 个有效字段定义")
    if error_lines:
        print(f"发现 {len(error_lines)} 个错误行")
    
    # 分析字段类型
    type_stats = analyze_column_types(columns)
    print("\n数据类型分布:")
    for data_type, count in sorted(type_stats.items()):
        print(f"  {data_type}: {count} 个字段")
    
    # 生成批量ALTER TABLE语句
    statements = generate_batch_alter_statements(columns, target_table_name, full_table_name)
    
    if not statements:
        print("没有生成任何SQL语句，程序退出")
        return
    
    print(f"\n生成了 {len([s for s in statements if s.strip() and not s.strip().startswith('--')])} 条SQL语句")
    
    # 保存到文件
    save_statements_with_summary(statements, output_file, type_stats, error_lines)
    
    print("-" * 50)
    print("处理完成！")
    print("\n重要提醒:")
    print("1. 脚本使用事务包装，确保原子性操作")
    print("2. 执行前请在测试环境验证")
    print("3. 如果字段已存在，会导致执行失败")
    print("4. 建议分批执行，避免长时间锁表")
    print("5. 执行前请备份数据库")

if __name__ == '__main__':
    main()